import os
from pdbench.utils.sampler import BaseSampler


class AlphaFlowSampler(BaseSampler):
    def __init__(self,
                 fasta_file_path: str,
                 sample_size: int,
                 save_path: str):
        super().__init__(fasta_file_path,
                         sample_size,
                         save_path)

    def sample_single(self,
                      k: str,
                      sequence: str
                      ):
        sample_save_path = os.path.join(self.save_path, k.split(" ")[0])
        os.mkdir(sample_save_path)

